#include <fstream>
#include <iomanip>
#include <sstream>
#include <sys/resource.h>

#include "csrMatrix.hpp"
/*! krylov */
#include "DILUPrecond.hpp"
#include "DiagPrecond.hpp"
#include "MGPrecond.hpp"
#include "PBiCGStab.hpp"

/*! AMG */
#include "aggregation.hpp"
#include "unapMultigrid.hpp"

static inline void mem_usage(long long int&total) {
  struct rusage r_usage;
  getrusage(RUSAGE_SELF, &r_usage);
  long long int mymem = r_usage.ru_maxrss;
  MPI_Allreduce(&mymem, &total, 1, MPI_LONG_LONG_INT, MPI_SUM, MPI_COMM_WORLD);
  //int myrank;
  //MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
  //if (myrank == 0) {
  //fprintf(stdout, "Total memory usage: %.2f MB\n", (double)((double)(total)) / (double)(1024));
  //fflush(stdout);  
  //}
}
/*! 定位文件绝对路径相关宏 */
#define LOCATEMATRIX(newName, pre, suff, dir)                         \
  {                                                                   \
    std::ostringstream os;                                            \
    os << dir << pre << (pre.empty() ? "" : "_") << "matrix_" << suff \
       << ".dat";                                                     \
    strcpy(newName, os.str().c_str());                                \
  }
#define LOCATEX(newName, pre, suff, dir)                                   \
  {                                                                        \
    std::ostringstream os;                                                 \
    os << dir << pre << (pre.empty() ? "" : "_") << "vector_init_" << suff \
       << ".dat";                                                          \
    strcpy(newName, os.str().c_str());                                     \
  }
#define LOCATEB(newName, pre, suff, dir)                                  \
  {                                                                       \
    std::ostringstream os;                                                \
    os << dir << pre << (pre.empty() ? "" : "_") << "vector_rhs_" << suff \
       << ".dat";                                                         \
    strcpy(newName, os.str().c_str());                                    \
  }
/*! * * * * * * * * * * * * * * *  * * * * * * * * * * * * * */
using namespace UNAP;
/*! 用于排序的辅助函数 */
template <typename Cmpt>
inline static bool compareTri(const Triplet<Cmpt>& a, const Triplet<Cmpt>& b) {
  if (a.r_ != b.r_) return (a.r_ < b.r_);
  return (a.c_ < b.c_);
};

/*! 确定文件行数 */
int getLineNumber(const char* filename) {
  FILE* fp = fopen(filename, "r");
  label line = 0;
  const int maxCline = 256;
  char cline[maxCline];
  while (fgets(cline, maxCline, fp) != NULL) line++;
  fclose(fp);
  return line;
}
/*! 确定矩阵按进程划分 */
void getMyrange(label& myBeginInd, label& myEndInd, label tsize, int myId,
                int nProcs) {
  label avgColNum = tsize / nProcs;
  label remainder = tsize % nProcs;
  if (myId < remainder) {
    myBeginInd = myId * (avgColNum + 1);
    myEndInd = (myId + 1) * (avgColNum + 1);
  } else {
    myBeginInd = myId * avgColNum + remainder;
    myEndInd = (myId + 1) * avgColNum + remainder;
  }
}
/*! 读取文件构建向量 */
void buildVecFromFile(scalarVector& vec, label myBeginInd, label myEndInd,
                      const char* filename) {
  FILE* fp = fopen(filename, "r");
  if (fp == NULL) {
    std::cerr << "ERROR: could not read file";
    exit(1);
  }
  const int max_cline = 256;
  char cline[max_cline];
  if (fgets(cline, max_cline, fp) == NULL) { /*!< skip first line */
    std::cerr << "ERROR: could not read from file" << std::endl;
    exit(1);
  }
  scalar* data = vec.data();
  scalar bData;
  label count = 0;
  while (fscanf(fp, "%lf\n", &bData) != EOF) {
    if (count >= myBeginInd && count < myEndInd) {
      data[count - myBeginInd] = bData;
    }
    count++;
    if (count >= myEndInd) break;
  }
  fclose(fp);
}

void buildVecFromFile2(scalarVector& vec, label myBeginInd, label myEndInd,
                      const char* filename) {
  FILE* fp = fopen(filename, "r");
  if (fp == NULL) {
    std::cerr << "ERROR: could not read file";
    exit(1);
  }
  const int max_cline = 256;
  char cline[max_cline];
  if (fgets(cline, max_cline, fp) == NULL) { /*!< skip first line */
    std::cerr << "ERROR: could not read from file" << std::endl;
    exit(1);
  }
  scalar* data = vec.data();
  scalar bData,x,y,z;
  label ind,count = 0;
  while (fscanf(fp, "%d %lf %lf %lf %lf \n", &ind,&bData,&x,&y,&z) != EOF) {
    if (count >= myBeginInd && count < myEndInd) {
      data[count - myBeginInd] = bData;
    }
    count++;
    if (count >= myEndInd) break;
  }
  fclose(fp);
}

/*! 向量输出工具 */
#define FILEOPEN(fout, fileName, MYID)     \
  std::ostringstream os;                   \
  os << fileName << "_" << MYID << ".txt"; \
  fout.open(os.str().c_str())
#define FILECLOSE(fout, fileName) fout.close()
template <typename T>
void printVector(const T& v, const char* fileName) {
  const label size = v.size();
  std::ofstream fout;

  FILEOPEN(fout, fileName, v.getCommunicator()->getMyId());

  fout << "rowSize: " << size << std::endl;
  forAll(i, size) {
    fout << std::setiosflags(std::ios::scientific) << std::setprecision(15)
         << v[i] << std::endl;
  }
  FILECLOSE(fout, fileName);
}

/*! * * * * * * * * * * * * * * * 主函数  * * * * * * * * * * * * * */
int main(int argc, char* argv[]) {
  COMM::createLogFile(false);
  /* Initialize MPI */
  initUtility();
  long long int total0,total1;
  mem_usage(total0);
  Communicator* comm = &COMM::getGlobalComm();
  const int nProcs = comm->getMySize();
  const int myId = comm->getMyId();
  label gSize, size, gNnz, nnz, beginR, endR;

  char dir[200];
  char fileName[200];
  string field;
  if (argc > 1)
    strcpy(dir, argv[1]);
  else
    strcpy(dir, "../kvlcc2/");
  if (argc > 2)
    field = argv[2];
  else
    field = "u";

  string pre = "";
  string method;
  if (argc > 3)
    method = argv[3];
  else
    method = "Krylov";

  string suff = field + "_101";
  /*!< get basic info */
  LOCATEB(fileName, pre, suff, dir)
  gSize = getLineNumber(fileName) - 2;
  getMyrange(beginR, endR, gSize, myId, nProcs);
  size = endR - beginR;
  scalarVector x(size, comm), b(size, comm);

  /*!< read vector */
  buildVecFromFile(b, beginR, endR, fileName);
  LOCATEX(fileName, pre, suff, dir)
  buildVecFromFile2(x, beginR, endR, fileName);

  /*! read matrix */
  LOCATEMATRIX(fileName, pre, suff, dir);
  gNnz = getLineNumber(fileName);
  CSRMatrix<scalar> A(comm);
  A.rows() = gSize;
  A.nnz() = gNnz;
  readMatrixMarket(A, fileName);
  //printVector(x, "x0");
  //printVector(b, "b");
  /*! * * * * * * * * * * * * * * * 开始计算 * * * * * * * * * * * * * */
  comm->barrier();
  double start_CPU, end_CPU;
  start_CPU = MPI_Wtime();  // 计时函数

  /*! * * * * * * * * * * * * * * * 1  * * * * * * * * * * * * * */
  SolverPerformance solverPerf;
  if (method == "Krylov") {
    DiagPrecond<scalar> precond(A);
    PBiCGStab<scalar> PBiCGStab(precond);

    PBiCGStab.SET_minIter(1);
    PBiCGStab.SET_maxIter(100);
    PBiCGStab.SET_ifPrint(false);
    PBiCGStab.SET_relTol(1e-6);

    solverPerf = PBiCGStab.solve(x, A, b);
  } else if(method =="Krylov2"){
    DILUPrecond<scalar> precond(A);
    PBiCGStab<scalar> PBiCGStab(precond);

    PBiCGStab.SET_minIter(1);
    PBiCGStab.SET_maxIter(100);
    PBiCGStab.SET_ifPrint(false);
    PBiCGStab.SET_relTol(1e-6);

    solverPerf = PBiCGStab.solve(x, A, b);
  } else {
    /*! * * * * * * * * * * * * * * * 2  * * * * * * * * * * * * * */

    Aggregation<scalar> aggl(A);
    aggl.SET_maxLevels(50);
    aggl.SET_rowSizeInCoarsestLevel(50); 
    // AMG setup
    aggl.agglomerate(); 

    MGSolver<scalar> MG(A, aggl);

    // set multigrid controls
    // not work when as a preconditioner
    MG.SET_tolerance(0.);
    MG.SET_relTol(1e-6);
    MG.SET_maxIter(100);
    MG.SET_ifPrint(false);

    MG.SET_nPreSweeps(1);
    MG.SET_nPostSweeps(1);
    MG.SET_maxPreSweeps(2);
    MG.SET_maxPostSweeps(2);
    MG.SET_smootherType("Jacobi");
    // -1 : V-cycle
    MG.SET_KcycleLevel(-1);
    MG.SET_scaleCorrection(false);

    MG.initSmoothers();
    MG.setUp();
    MG.printStat(false);

    MGPrecond<scalar> precond(A, &MG);
    PBiCGStab<scalar> PBiCGStabSolver(precond);

    PBiCGStabSolver.SET_minIter(1);
    PBiCGStabSolver.SET_maxIter(100);
    PBiCGStabSolver.SET_ifPrint(false);
    PBiCGStabSolver.SET_relTol(1e-6);

    solverPerf = PBiCGStabSolver.solve(x, A, b);
    //solverPerf = MG.solve(x, A, b);
  }

  /*! * * * * * * * * * * * * * * *  * * * * * * * * * * * * * */

  end_CPU = MPI_Wtime();  // 计时函数


  double time_CPU = ((double)end_CPU - (double)start_CPU);
  MPI_Allreduce(MPI_IN_PLACE, &time_CPU, 1, MPI_DOUBLE, MPI_MAX,
                comm->getMpiComm());
  // output to file
  //printVector(x, field.c_str());

  COUT << "\nAfter " << solverPerf.nIterations()
       << " iterations, the solution is converged!" << ENDL
       << "final relative Res "
       << solverPerf.finalResidual() / solverPerf.initialResidual() << ENDL;
  if (myId == 0) {
    printf("max_time_CPU=%f\n", time_CPU);
  }

  mem_usage(total1);
  COUT<<"Memory usage: "<<(double)((double)(total1-total0)) / (double)(1024) <<" MB\n";
} // end of main
